# 取一个 reflow 数据集 ds
# 从 ds 中取一个样例 data
# 取得 data 的 noise 和 latent
# 均匀采样不同时间点的 t
# 计算不同 t 之下的 loss
# 将所有的 loss 保存，打印出图像

# %%
# 确定配置
import torch
from reflow.utils import _PIPELINES, _SCHEDULERS, set_seed
from diffusers import UNet2DConditionModel
from reflow.utils import decode_latents
from tqdm.auto import tqdm, trange
from torchvision.utils import make_grid
from typing import Callable, List, Optional, Union
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import math

gpu = 2
device = f'cuda:{gpu}' if gpu >= 0 else 'cpu'
diffusers_pipeline = 'stable_diffusion'
diffusers_scheduler = 'dpm_solver_multi'
diffusers_pipeline_ckpt = 'checkpoints/SD-1-4'
use_xformers = True
weight_dtype = torch.float16

# %%加载 diffusion pipeline;
pipeline_cls = _PIPELINES[diffusers_pipeline]
scheduler_cls = _SCHEDULERS[diffusers_scheduler]

weight_dtype=torch.float16
pipeline = pipeline_cls.from_pretrained(
    diffusers_pipeline_ckpt,
    torch_dtype=weight_dtype,
    safety_checker=None,
    requires_safety_checker=False,
)
pipeline.scheduler = scheduler_cls.from_config(pipeline.scheduler.config)
pipeline.vae.load_state_dict(torch.load("checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.bin"))
pipeline = pipeline.to(device)

if use_xformers:
    pipeline.enable_xformers_memory_efficient_attention()

# %%
# 设定不同的 t compressed reflow model , 具有不同的 ckpt path 和 zt_compress_rate

reflow_model_ckpt_paths = [
    'logs/zt_compressed/laion_100kPrompts_guidance1.0_zt0.2/checkpoints/score_model_s1600.pth',
    'logs/zt_compressed/laion_100kPrompts_guidance1.0_zt0.15/checkpoints/score_model_s1600.pth',
    'logs/zt_compressed/laion_100kPrompts_guidance1.0_zt0.1/checkpoints/score_model_s1600.pth',
    'logs/zt_compressed/laion_100kPrompts_guidance1.0_zt0.05/checkpoints/score_model_s1600.pth',
]
zt_compress_rates = [0.2, 0.15, 0.1, 0.05]

# ! debug
reflow_model_ckpt_paths = []
zt_compress_rates = []

# 添加 base reflow model 的路径
reflow_model_ckpt_paths.append(
    'logs/online/sdv1-4_laion_1MPrompts_guidance7.5/checkpoints/score_model_s40000.pth'
)


def load_reflow_model(reflow_model_ckpt_path):
    unet = UNet2DConditionModel.from_pretrained(
        diffusers_pipeline_ckpt,
        subfolder='unet',
    )
    unet.load_state_dict(torch.load(reflow_model_ckpt_path,
                         map_location='cpu'), strict=True)
    unet.to(device, dtype=weight_dtype).eval().requires_grad_(False)
    if use_xformers:
        unet.enable_xformers_memory_efficient_attention()
    return unet


class HybridModel(torch.nn.Module):
    def __init__(self, zt_compressed_model, base_model, zt_compress_rate) -> None:
        super().__init__()
        self.zt_compressed_model = zt_compressed_model
        self.base_model = base_model

        self.zt_compress_rate = zt_compress_rate
        self.eps = 1e-3

    def forward(self, sample, timestep, **condition):
        if timestep < (999*(self.zt_compress_rate+self.eps)):
            timestep = timestep / zt_compress_rate
            model = self.zt_compressed_model
        else:
            model = self.base_model
        return model(sample, timestep=timestep, **condition)

    @property
    def device(self):
        return self.base_model.device


# %%# 损失函数


def loss_fn(score, target, type='l2'):
    if type == 'l2':
        loss = (score-target).square()
    elif type == 'l1':
        loss = (score-target).abs()
    return loss

# %%# loss curve over t (pred=model(zt, t, c))



def zt_loss(model, timesteps, z0, z1, condition,):
    loss_list = []
    target = z1 - z0
    with torch.no_grad():
        for t in tqdm(timesteps[:-1]):
            # TODO
            t_expand = t.view(-1, 1, 1, 1)
            zt = t_expand * z1 + (1.-t_expand) * z0

            t = (999*t)
            score = model(zt, timestep=t, **condition).sample

            loss = loss_fn(score, target)
            loss_list.append(loss.mean().detach().cpu().item())
    zt_loss_list = np.array(loss_list)

    return zt_loss_list
# %%# 边采样边计算 loss (pred=model(sample, t, c)) 但是可以 cheat (前 x 步使用 answer) (not in use)



def sample_loss(model, timesteps, z0, z1, condition, uncond_condition=None, guidance_scale=1.0, no_guide_t: float = 0,):

    loss_list = []
    sample = z0
    target = z1 - z0

    with torch.no_grad():
        for i in trange(len(timesteps)-1):
            t = timesteps[i]
            t_n = timesteps[i+1]
            vec_t = (999*t)
            pred = model(sample, timestep=vec_t, **condition).sample
            loss = loss_fn(pred, target)

            loss_list.append(loss.mean().detach().cpu().item())

            dt = t_n - t
            if uncond_condition and guidance_scale > 1.0 and t > no_guide_t:
                uncond_pred = model(sample, timestep=vec_t,
                                    **uncond_condition).sample
                pred = uncond_pred + guidance_scale * (pred - uncond_pred)
            sample = sample + pred * dt
    sample_loss_list = np.array(loss_list)

    return sample_loss_list, sample


# %%采样设置, 加载 prompts, 定义 noise
random_seed = 0
if random_seed != -1:
    set_seed(random_seed)

persistent_noise = True
noise = None

prompts = "A young man smiles and holds a small teddy bear."
if not (persistent_noise and isinstance(noise, torch.Tensor)):
    noise = torch.randn(1, 4, 64, 64, device=device, dtype=weight_dtype)

# %%使用 diffusion model 采样获取数据
inference_steps = 25
guidance_scale = 7.5

example = pipeline.inference_latent(
    prompt=prompts,
    latents=noise, 
    num_inference_steps=inference_steps,
    guidance_scale=guidance_scale,
)
# noise = example['noise']
latent = example['latent']
condition = {
    'encoder_hidden_states': example['text_embeddings']
}
uncond_condition = {
    'encoder_hidden_states': example['uncond_text_embeddings']
}
z0 = noise.to(dtype=weight_dtype)
z1 = latent.to(dtype=weight_dtype)

# %%# 使用 uniform reflow model 绘制 sample loss 曲线; 加载其余的 t compressed model , 绘制 loss 曲线
N = 1
guidance_scale = 1.5

sample_loss_curves = []
samples = []
ts_list = []

base_reflow_model = load_reflow_model(
    reflow_model_ckpt_paths[-1],
)
Ts, Te = 1e-3, 1
timesteps = torch.arange(start=Ts, end=Te, step=(Te-Ts)/N)
timesteps = torch.tensor(timesteps.tolist() +
                         [Te], device=device, dtype=weight_dtype)

ts_list.append(timesteps[:-1].cpu().numpy())
sample_loss_list, sample = sample_loss(base_reflow_model, timesteps, z0, z1,
                                       condition, uncond_condition=uncond_condition, guidance_scale=guidance_scale)
sample_loss_curves.append(sample_loss_list)
samples.append(sample)

for i in range(len(zt_compress_rates)):
    zt_compress_rate = zt_compress_rates[i]
    reflow_model_ckpt_path = reflow_model_ckpt_paths[i]
    ts_list.append(timesteps[:-1].cpu().numpy())

    hybrid_model = HybridModel(
        load_reflow_model(reflow_model_ckpt_path),
        base_reflow_model,
        zt_compress_rate,
    )

    sample_loss_list, sample = sample_loss(hybrid_model, timesteps, z0, z1, condition,
                                           uncond_condition=uncond_condition, guidance_scale=guidance_scale, no_guide_t=0)
    sample_loss_curves.append(sample_loss_list)
    samples.append(sample)

# %%# plot sample loss over t

plt.plot(ts_list[0], sample_loss_curves[0], label='base')
for i in range(len(zt_compress_rates)):
    zt_compress_rate = zt_compress_rates[i]
    plt.plot(ts_list[i+1], sample_loss_curves[i+1],
             label=f'zt_compress_{zt_compress_rate}')

plt.title(f'{prompts} (N_diffusion={inference_steps} N_base={N} N_hybrid={N})')
plt.xlabel("t")
plt.ylabel("loss")
plt.legend()
plt.show()

# %%# 打印 images

images = [decode_latents(pipeline.vae, z1).squeeze()]
for sample in samples:
    images.append(decode_latents(pipeline.vae, sample).squeeze())
nrow = int(math.sqrt(len(images)))
nrow = max(nrow, len(images)//nrow)
grid = make_grid(
    images,
    nrow=nrow,
    pad_value=0,
)
grid = grid.mul(255.).to(dtype=torch.uint8)
grid = grid.permute(1, 2, 0).numpy()
pil_grid = Image.fromarray(grid)

print('(1,1) latent\n(1,2) base')
cnt = 2
for zcr in zt_compress_rates:
    print(f'({cnt//nrow+1},{cnt%nrow+1}) zt compress rate {zcr}')
    cnt += 1

plt.figure(figsize=(20, 20))
plt.title(f'{prompts} (N_diffusion={inference_steps} N_base={N} N_hybrid={N})')
plt.imshow(pil_grid)

# %%# 使用 uniform reflow model 绘制 zt loss 曲线; 加载其余的 t compressed model , 绘制 loss 曲线


zt_loss_curves = []
ts_list = []

base_reflow_model = load_reflow_model(
    reflow_model_ckpt_paths[-1],
)
N = 100
Ts, Te = 1e-3, 1
timesteps = torch.arange(start=Ts, end=Te, step=(Te-Ts)/N)
timesteps = torch.tensor(timesteps.tolist() +
                         [Te], device=device, dtype=weight_dtype)

ts_list.append(timesteps[:-1].cpu().numpy())
zt_loss_curves.append(zt_loss(base_reflow_model, timesteps, z0, z1, condition))

for i in range(len(zt_compress_rates)):
    zt_compress_rate = zt_compress_rates[i]
    reflow_model_ckpt_path = reflow_model_ckpt_paths[i]

    N1, N2 = 25, 25
    Ts, Te = 1e-3, 1
    timesteps1 = torch.arange(
        start=Ts, end=zt_compress_rate, step=(zt_compress_rate-Ts)/N1)
    timesteps2 = torch.arange(start=zt_compress_rate, end=Te, step=(
        Te-zt_compress_rate)/N2)
    timesteps = torch.cat([timesteps1, timesteps2], dim=0)
    timesteps = torch.tensor(timesteps.tolist() +
                             [Te], device=device, dtype=weight_dtype)
    ts_list.append(timesteps[:-1].cpu().numpy())

    hybrid_model = HybridModel(
        load_reflow_model(reflow_model_ckpt_path),
        base_reflow_model,
        zt_compress_rate,
    )

    zt_loss_curves.append(zt_loss(hybrid_model, timesteps, z0, z1, condition))
# %%# plot zt loss over t

plt.plot(ts_list[0], zt_loss_curves[0], label='base')
for i in range(len(zt_compress_rates)):
    zt_compress_rate = zt_compress_rates[i]
    plt.plot(ts_list[i+1], zt_loss_curves[i+1],
             label=f'zt_compress_{zt_compress_rate}')

plt.title(f'{prompts}')
plt.xlabel("t")
plt.ylabel("loss")
plt.legend()
plt.show()

# %%
